/*
 Copyright (c) 2023 Broin tianjianchen All rights reserved.
 Use of this source code is governed by a BSD-style
 license that can be found in the LICENSE file.
 */

#include <argp.h>
#include <stdio.h>
#include <unistd.h>
#include <bpf/libbpf.h>
#include <sys/resource.h>
#include <time.h>
#include <sys/stat.h>

#include "lsm_permission.skel.h"

struct lsm_permission_env
{
    int verbose;
    int is_csv;
    pid_t target_pid;
    pid_t exclude_current_ppid;
    long min_duration_ms;
    volatile int *exiting;
};

struct ctx_data
{
    int whole_file_policy_map_fd;
    int file_policy_map_fd;
};

struct abs_con
{
    int container_policy_map_fd;
    int whole_file_policy_map_fd;
};

static int audit_output(void *ctx, void *data, size_t data_sz)
{
    const struct audit_event *e = data;
    if (AUDIT_TYPE_FILE == e->type)
    {
        if (LSM_ALLOW == e->decision)
        {
            ebpflsm_info("access passed, ino: %llu, dev: %lu, comm: %s, access: %lu, pid: %lu, tgid: %lu",
                         e->file_audit.st_ino, e->file_audit.st_dev, e->comm, e->file_audit.access, e->pid, e->tgid);
        }
        else if (LSM_DENY == e->decision)
        {
            ebpflsm_info("access denied, ino: %llu, dev: %lu, comm: %s, access: %lu, pid: %lu, tgid: %lu",
                         e->file_audit.st_ino, e->file_audit.st_dev, e->comm, e->file_audit.access, e->pid, e->tgid);
        }
        else
        {
            ebpflsm_normal("access not care, ino: %llu, dev: %lu, comm: %s, access: %lu, pid: %lu, tgid: %lu",
                           e->file_audit.st_ino, e->file_audit.st_dev, e->comm, e->file_audit.access, e->pid, e->tgid);
        }
    }

    if (AUDIT_TYPE_NET == e->type)
    {
        if (LSM_ALLOW == e->decision)
        {
            ebpflsm_info("access passed,ip: %s, uid: %lu, comm: %s, access: %lu ,pid: %lu, tgid: %lu",
                         convert_ipbin(e->net_audit.ip), e->net_audit.uid, e->comm, e->net_audit.opt, e->pid, e->tgid);
        }
        else if (LSM_DENY == e->decision)
        {
            ebpflsm_info("access denied,ip: %s, uid: %lu, comm: %s, access: %lu ,pid: %lu, tgid: %lu",
                         convert_ipbin(e->net_audit.ip), e->net_audit.uid, e->comm, e->net_audit.opt, e->pid, e->tgid);
        }
        else
        {
            ebpflsm_info("access not care,ip: %s, uid: %lu, comm: %s, access: %lu ,pid: %lu, tgid: %lu",
                         convert_ipbin(e->net_audit.ip), e->net_audit.uid, e->comm, e->net_audit.opt, e->pid, e->tgid);
        }
    }

    return 0;
}

static int old_inode_update(void *ctx, void *data, size_t data_sz)
{
    const struct mv_old_inode *event = (struct mv_old_inode *)data;
    struct ctx_data *ctx_info = (struct ctx_data *)ctx;
    struct file_policy_val value;
    char proc_cwd[PATH_MAX];
    char abs_new_path[PATH_MAX];
    struct stat st;
    printf("new path is %s\n", event->new_path);
    snprintf(proc_cwd, sizeof(proc_cwd), "/proc/%d/cwd", event->pid);

    // 检查进程是否仍然存在
    if (kill(event->pid, 0) == -1 && errno == ESRCH)
    {
        fprintf(stderr, "进程 %d 不存在: %s\n", event->pid, strerror(errno));
        return 0; // 忽略这个错误，继续执行
    }

    // 解析新路径为绝对路径
    if (event->new_path[0] != '/')
    {
        ssize_t len_new = readlink(proc_cwd, abs_new_path, sizeof(abs_new_path) - 1);
        if (len_new != -1)
        {
            abs_new_path[len_new] = '\0';
            snprintf(abs_new_path + len_new, sizeof(abs_new_path) - len_new, "/%s", event->new_path);
        }
        else
        {
            fprintf(stderr, "获取 PID %d 的 CWD 失败: %s\n", event->pid, strerror(errno));
            return 0; // 忽略这个错误，继续执行
        }
    }
    else
    {
        snprintf(abs_new_path, sizeof(abs_new_path), "%s", event->new_path);
    }

    // 获取新路径的 inode 号
    if (stat(abs_new_path, &st) == 0)
    {

        // 如果前后inode号不相同
        if (event->inode_id != st.st_ino)
        {
            int whole_file_policy_map_fd = ctx_info->whole_file_policy_map_fd;
            // struct file_policy_val whole_file = {};
            int ret = bpf_map_lookup_elem(whole_file_policy_map_fd, &event->inode_id, &value);
            if (!ret)
            {
                // 更新map
                ret = bpf_map_update_elem(whole_file_policy_map_fd, &st.st_ino, &value, BPF_ANY);
                if (ret)
                {
                    fprintf(stderr, "添加new_inode失败");
                    return 0;
                }
                ret = bpf_map_delete_elem(whole_file_policy_map_fd, &event->inode_id);
                if (ret)
                {
                    fprintf(stderr, "删除old_inode失败");
                    return 0;
                }
            }
            else
                return 0;
        }
        else
        {
            printf("前后inode相同\n");
        }
    }
    else
    {
        fprintf(stderr, "获取新路径 stat 信息失败: %s\n", abs_new_path);
        // 找不到新路径
    }
    return 0;
}

/*容器获取绝对路径*/
static int container_abs_path(void *ctx, void *data, size_t data_sz)
{
    const struct container_event *event = data;
    // printf("current pid is %d, mntns is %lu\n", event->pid, event->mntns);
    struct abs_con *ctx_info = (struct abs_con *)ctx;
    int container_policy_map_fd = ctx_info->container_policy_map_fd;
    int whole_file_policy_map_fd = ctx_info->whole_file_policy_map_fd;
    char merged_path[256];
    int pid = event->pid;
    // unsigned long mntns = event->mntns;

    /*这一部分是获取绝对路径*/
    char mountinfo_path[256];
    snprintf(mountinfo_path, sizeof(mountinfo_path), "/proc/%d/mountinfo", pid);

    // printf("mountinfo_path is %s\n",mountinfo_path);
    FILE *fp = fopen(mountinfo_path, "r");
    if (!fp)
    {
        ebpflsm_error("can't open %d\n", pid);
        return 0;
    }
    char line[1024];
    while (fgets(line, sizeof(line), fp))
    {
        if (strstr(line, "workdir="))
        {
            char *workdir = strstr(line, "workdir=") + 8;
            char *end = strchr(workdir, ',');
            if (end)
                *end = '\0';
            // printf("Workdir: %s\n", workdir);

            strncpy(merged_path, workdir, sizeof(merged_path));

            char *work_pos = strstr(merged_path, "/work");
            if (work_pos)
            {

                memcpy(work_pos, "/merged", 7);
                work_pos[7] = '\0';
            }

            // printf("Merged path: %s\n", merged_path);
            break;
        }
    }
    fclose(fp);

    /*读取容器的镜像名称*/
    char cmd[256];
    snprintf(cmd, sizeof(cmd), "docker ps -q");

    fp = popen(cmd, "r");
    if (fp == NULL)
    {
        fprintf(stderr, "Failed to run command: %s\\n", strerror(errno));
        return 0;
    }

    char container_id[64];
    char result[256] = {0};
    int found = 0;

    while (fgets(container_id, sizeof(container_id), fp) != NULL)
    {
        // 去掉末尾的换行符
        container_id[strcspn(container_id, "\\n")] = 0;

        snprintf(cmd, sizeof(cmd), "docker inspect --format '{{.State.Pid}}' %s", container_id);
        FILE *fp_pid = popen(cmd, "r");
        if (fp_pid == NULL)
        {
            fprintf(stderr, "Failed to run command: %s\\n", strerror(errno));
            continue;
        }

        char pid_str[16];
        if (fgets(pid_str, sizeof(pid_str), fp_pid) != NULL)
        {
            int container_pid = atoi(pid_str);
            if (container_pid == event->pid)
            {
                snprintf(cmd, sizeof(cmd), "docker inspect --format '{{.Config.Image}}' %s", container_id);
                FILE *fp_image = popen(cmd, "r");
                if (fp_image == NULL)
                {
                    fprintf(stderr, "Failed to run command: %s\\n", strerror(errno));
                    continue;
                }

                if (fgets(result, sizeof(result), fp_image) != NULL)
                {
                    // printf("Container Image: %s", result);
                    found = 1;
                }
                pclose(fp_image);
                break;
            }
        }
        pclose(fp_pid);
    }

    if (!found)
    {
        ebpflsm_error("No container found for PID %u\n", event->pid);
    }
    else
    {
        char *newline = strchr(result, '\n');
        if (newline)
        {
            *newline = '\0';
        }
        policy_container_val_t policy_container_val;
        // printf("result is %s\n",result);
        int ret = bpf_map_lookup_elem(container_policy_map_fd, result, &policy_container_val);
        if (!ret)
        {
            struct stat st;
            strcat(merged_path, policy_container_val.path);
            ebpflsm_info("find %s policy ,merge path is %s\n", result, merged_path);
            if (stat(merged_path, &st) == 0)
            {
                long long ino = st.st_ino;
                file_policy_val_t policy_val = {0x0, 0x0};
                ret = bpf_map_lookup_elem(whole_file_policy_map_fd, &ino, &policy_val);
                if (ret)
                {
                    policy_val.allow |= policy_container_val.allow;
                    policy_val.deny |= policy_container_val.deny;
                    ret = bpf_map_update_elem(whole_file_policy_map_fd, &ino, &policy_val, BPF_ANY);
                }
            }
        }
        else
        {
            printf("not found\n");
        }
    }
    pclose(fp);

    return 0;
}

/**
 * @description:
 * @param {struct lsm_permission_env} env:全局过滤条件设置
 * @param {struct lsm_permission_bpf *} skel: CO-RE骨架，承载内核态相关信息
 * @param {void *} ctx: 用户态相关上下文信息，用以推送到handle_event进行进一步的条件过滤和其他用途
 * @return {*} 不用管
 */

static int start_lsm_permission_tracker(struct lsm_permission_env env, struct lsm_permission_bpf *skel, void *ctx)
{
    int err;

    // 配置信息
    char *config_path = "../../config/config.yaml";

    // 审计
    struct ring_buffer *rb = NULL;

    struct ring_buffer *rb_old = NULL;

    // container
    struct ring_buffer *rb_container = NULL;

    /* 初始化开始 */

    // 配置相关信息
    err = config_setup(config_path);

    if (!env.exiting)
    {
        fprintf(stderr, "env.exiting is not set.\n");
        return -1;
    }
    libbpf_set_strict_mode(LIBBPF_STRICT_ALL);

    skel = lsm_permission_bpf__open();
    if (!skel)
    {
        fprintf(stderr, "Failed to open and load BPF skeleton\n");
        return 1;
    }

    /* 加载并验证BPF程序 */
    err = lsm_permission_bpf__load(skel);
    if (err)
    {
        fprintf(stderr, "Failed to load and verify BPF skeleton.\n");
        goto cleanup;
    }

    /* 挂载监控点 */
    err = lsm_permission_bpf__attach(skel);
    if (err)
    {
        fprintf(stderr, "Failed to attach BPF skeleton.\n");
        goto cleanup;
    }

    /* 初始化结束 */

    struct ctx_data ctx_info = {
        .whole_file_policy_map_fd = bpf_map__fd(skel->maps.whole_file_policy),
        .file_policy_map_fd = bpf_map__fd(skel->maps.file_policy) // 假设还有另一个 map
    };

    struct abs_con ctx_container = {
        .container_policy_map_fd = bpf_map__fd(skel->maps.container_policy),
        .whole_file_policy_map_fd = bpf_map__fd(skel->maps.whole_file_policy)
    };

    // 通过策略文件引入策略
    ROLE_MAP *role_map = hash_map_create(MAX_ROLE_SIZE);
    policy_add_all_file_to_map(config.policy_path, role_map, skel);

    /* 审计输出 */
    rb = ring_buffer__new(bpf_map__fd(skel->maps.event_audit), audit_output, ctx, NULL);
    if (!rb)
    {
        err = -1;
        ebpflsm_error("Failed to create ring buffer!");
        fprintf(stderr, "Failed to create ring buffer!\n");
        goto cleanup;
    }

    /* 旧inode_map */
    rb_old = ring_buffer__new(bpf_map__fd(skel->maps.old_pid_inode), old_inode_update, &ctx_info, NULL);
    if (!rb_old)
    {
        err = -1;
        ebpflsm_error("Failed to update old inode!");
        fprintf(stderr, "Failed to create ring buffer!\n");
        goto cleanup;
    }

    /*container获取绝对路径*/
    rb_container = ring_buffer__new(bpf_map__fd(skel->maps.container_path), container_abs_path, &ctx_container, NULL);
    if (!rb_container)
    {
        err = -1;
        ebpflsm_error("Failed to get abs path of container!");
        fprintf(stderr, "Failed to create ring buffer!\n");
        goto cleanup;
    }

    while (!*env.exiting)
    {
        err = ring_buffer__poll(rb, 100);
        if (err < 0 && err != -EINTR)
        {
            fprintf(stderr, "Error polling rb buffer: %d\n", err);
            break;
        }
        err = ring_buffer__poll(rb_old, 100);
        if (err < 0 && err != -EINTR)
        {
            fprintf(stderr, "Error polling rb buffer: %d\n", err);
            break;
        }
        err = ring_buffer__poll(rb_container, 100);
        if (err < 0 && err != -EINTR)
        {
            fprintf(stderr, "Error polling rb buffer: %d\n", err);
            break;
        }
    }

cleanup:
    hash_map_free(role_map);
    ring_buffer__free(rb);
    ring_buffer__free(rb_old);
    ring_buffer__free(rb_container);
    lsm_permission_bpf__destroy(skel);
    return err < 0 ? -err : 0;
}
